import os
import csv
from zkl_aiutils_datasets import load_dataset
from zkl_pyutils_fsspec import FsLike, resolve_fs

class EmbeddingReader:
    def __init__(self, dataset_path):
        self.dataset = self.dataset_read_init(dataset_path)
        self.str_to_id = self.csv_init(dataset_path)

    
    def dataset_read_init(self, dataset_path):
        fs = resolve_fs(os.path.join(dataset_path, "dataset"))
        dataset = load_dataset(fs)
        return dataset

    def csv_init(self, dataset_path):
        str_to_id = {}
        with open(os.path.join(dataset_path, "str_list.csv"), 'r') as f:
            reader = csv.reader(f)
            next(reader)
            for i, row in enumerate(reader):
                str_to_id[row[0]] = i
        return str_to_id

    def get_embeddings_batch(self, texts):
        embeddings = []
        for text in texts:
            try:
                id = self.str_to_id[text]
                embeddings.append(self.dataset[id][0][0])
            except KeyError:
                # 如果在str_to_id中找不到对应的key，返回None
                embeddings.append(None)
        # print(embeddings)
        return embeddings
    
    def get_embedding(self, text):
        embedding = self.get_embeddings_batch([text])
        return embedding[0]
        
        
    def close(self):
        # self.dataset.close()
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def __del__(self):
        self.close()

if __name__ == "__main__":
    dataset_path = "/home/xw/python3test/omics-bert-datasets-2-embeddings/embeddings_SSD/ukbiobank/v4.2.8"
    embedding_reader = EmbeddingReader(dataset_path=dataset_path)
    embeddings = embedding_reader.get_embeddings_batch(["00000n0n0nn0n0n0n2n","0"])
    print(embeddings)